from lib.model import system_model
from itertools import product
import matplotlib.pyplot as plt
from tqdm import tqdm
import lib.tools as my_tools
import numpy as np
import time
import sys
import csv
import pickle

interception_system = system_model("interception_system")

def find_dataset():
    """
    对于方法step_interception_model, u必须是元胞数组，每个元胞对应一个参数，元胞的类型和动态模型的输入相同
    for example:
    u = (np.array([1]),np.array([0,0,0]),)
    initstate = np.array([0.5, 0.8, 100, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1])
    result = interception_system.step_interception_model(x0=initstate, input = u, run_time=0.003)
    print(f"result:\n{result.y}")
    """
    # 控制量v_zd
    v_zd  = np.arange(-2, 2, 0.2)
    # 控制量n_3d,希望飞机姿态前倾
    n_3d = np.array([0.2, 0, 0.98])
    # 控制量omega
    omega_x = np.arange(-0.2, 0.2, 0.05)
    omega_y = np.arange(-0.2, 0.2, 0.05)
    omega_z = np.arange(-0.8, 0.8, 0.1)


    # 不同图像误差
    p_x = np.arange(0.1, 0.4, 0.1)
    p_y = np.arange(0.1, 0.4, 0.1)
    # 不同物距
    c_z = np.arange(30, 50, 20)
    # 不同的前飞速度
    gv_z = np.arange(5, 10, 5)
    # 不同的n3向量，需要符合模为1的特点
    gen_n3 = interception_system.generate_n3()
    n3 = list(gen_n3)

    # 进度条参数
    print(f'p_x len {len(p_x)}')
    print(f'p_y len {len(p_y)}')
    print(f'c_z len {len(c_z)}')
    print(f'gv_z len {len(gv_z)}')
    print(f'n3 len {len(n3)}')
    state_len = len(p_x) * len(p_y) * len(c_z) * len(gv_z) * len(n3)
    state_pdt = product(range(len(p_x)), range(len(p_y)), range(len(c_z)), range(len(gv_z)), range(len(n3)))
    # 开始循环遍历
    data_item = []
    invalid_item = 0
    processing_item = 0
    my_tqdm = tqdm(state_pdt, total = state_len)
    for s1,s2,s3,s5,s6 in my_tqdm:
        data_item_sub = []
        input_pdt = product(range(len(v_zd)), range(len(omega_x)), range(len(omega_y)), range(len(omega_z)))
        has_positive_item = False
        for u1,u2,u3,u4 in input_pdt:

            # 计算D值
            Dvalue,D1,D2,D3 = interception_system.cal_D_value( 
                                            p_g = np.array([p_x[s1], p_y[s2]]), 
                                            c_z = c_z[s3], 
                                            gv_z = gv_z[s5], 
                                            n_3 = n3[s6],
                                            v_zd = v_zd[u1], 
                                            n_3d = n_3d,
                                            omega = np.array([omega_x[u2], omega_y[u3], omega_z[u4]]))
            
            # 保留所有D值为负的数据
            if Dvalue < 0:
                data_item_sub.append([Dvalue[0,0], D1[0,0], D2[0,0], D3[0,0], v_zd[u1], omega_x[u2], omega_y[u3], omega_z[u4], p_x[s1], p_y[s2], c_z[s3], gv_z[s5], n3[s6]])
            
            # 同时保留最小一条D值为正的数据
            if Dvalue >= 0:
                if not has_positive_item:
                    has_positive_item = True
                    data_item_positive = [Dvalue[0,0], D1[0,0], D2[0,0], D3[0,0], v_zd[u1], omega_x[u2], omega_y[u3], omega_z[u4], p_x[s1], p_y[s2], c_z[s3], gv_z[s5], n3[s6]]
                else:
                    if Dvalue[0,0] < data_item_positive[0]:
                        data_item_positive = [Dvalue[0,0], D1[0,0], D2[0,0], D3[0,0], v_zd[u1], omega_x[u2], omega_y[u3], omega_z[u4], p_x[s1], p_y[s2], c_z[s3], gv_z[s5], n3[s6]]

        # 将所有负的数据和最小的一个正的数据保存到data_item中
        data_item_sub.append(data_item_positive)
        data_item.append(data_item_sub)
        if len(data_item_sub) <= 1:
            invalid_item += 1
        processing_item +=1
        # 添加进度条前缀。未找到可用控制量的状态组/已经遍历过的状态组

        my_tqdm.set_description(f"invalid/processed item {invalid_item}/{processing_item},{invalid_item/processing_item*100:.2f}%")

    my_tools.save_dataset(data_item, 'found_data_set.pickle')

def find_dataset_n3():
    # 根据不同的状态n3, 和不同的期望n3d, 和omega_z找到合适omega_x omega_y, 使得最终的D<0
    
    # 控制量 omega
    # omega_x = np.arange(-0.1, 0.1, 0.005).astype(np.float32)
    omega_x = np.linspace(-0.1, 0.1, num=int((0.1 - (-0.1)) / 0.005)+1, dtype=np.float32)
    print(f'omega_x {omega_x}')
    omega_y = np.linspace(-0.1, 0.1, num=int((0.1 - (-0.1)) / 0.005)+1, dtype=np.float32)
    print(f'omega_y {omega_y}')

    data = [["omega_x len", len(omega_x)],
            ["omega_y len", len(omega_y)]]
    my_tools.print_table(data, "Control seed lengths")
    
    n_3d = []
    # n3: 对应倾斜姿态
    # n_3d_one = np.array([0.0, 0.2, 0.98], dtype=np.float32)
    # n_3d.append(n_3d_one)
    gen_n3d = interception_system.generate_n3(0.1)
    n_3d = list(gen_n3d)
    # state: omega_z 用于视觉伺服，所以这里看作是状态量
    omega_z = np.linspace(-0.8, 0.8, num=int((0.8 - (-0.8)) / 0.1)+1, dtype=np.float32)
    print(f'omega_z {omega_z}, type {omega_z.dtype}')
    # state: 不同的n3向量，需要符合模为1的特点
    gen_n3 = interception_system.generate_n3(0.1)
    n3 = list(gen_n3)
    # print(f"n3: {n3[1][0]}")

    data = [
        ["omega_z", len(omega_z)],
        ["n3", len(n3)],
        ["n_3d", len(n_3d)],
        ["all combinations", len(omega_z) * len(n3) * len(n_3d)]
    ]
    my_tools.print_table(data, "State seed lengths")

    # 根据状态组合数量添加带进度条的循环
    state_pdt = product(range(len(omega_z)), range(len(n3)), range(len(n_3d)))
    state_len = len(omega_z) * len(n3) * len(n_3d)
    state_tqdm = tqdm(state_pdt, total = state_len, ncols=110)
    # 开始循环遍历
    data_item = []
    invalid_item = 0
    processing_item = 0
    processing_item_in_chunk = 0
    n3_dataset_saver = my_tools.save_dataset_chunk(data_item, 'found_data_set_attitude.pickle')
    next(n3_dataset_saver)
    with open('state_invalid_attitude.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        # 写入表头
        writer.writerow(['omega_z', 'n3_x', 'n3_y', 'n3_z', 'n3_xd', 'n3_yd', 'n3_zd'])

        for s1, s2, s3 in state_tqdm:
            data_item_sub = []
            has_positive_item = False
            # 遍历当前状态组合下的所有解，无论稳定与否
            if np.array_equal(n3[s2], n_3d[s3]):
                # 如果当前状态和期望状态相同，则D值恒为0，跳过
                continue
            for u1, u2 in product(range(len(omega_x)), range(len(omega_y))):
                D_value = interception_system.cal_D_value_n3(n_3 = n3[s2], 
                                                            n_3d = n_3d[s3], 
                                                            omega = np.array([omega_x[u1], omega_y[u2], omega_z[s1]], dtype=np.float32)
                                                            )

                if D_value < 0:
                    data_item_sub.append([D_value[0,0], omega_x[u1], omega_y[u2], omega_z[s1], n3[s2][0], n3[s2][1], n3[s2][2], n_3d[s3][0], n_3d[s3][1], n_3d[s3][2]])

                if D_value >= 0:
                    if not has_positive_item:
                        has_positive_item = True
                        data_item_positive = [D_value[0,0], omega_x[u1], omega_y[u2], omega_z[s1], n3[s2][0], n3[s2][1], n3[s2][2], n_3d[s3][0], n_3d[s3][1], n_3d[s3][2]]
                    else:
                        if D_value[0,0] < data_item_positive[0]:
                            data_item_positive = [D_value[0,0], omega_x[u1], omega_y[u2], omega_z[s1], n3[s2][0], n3[s2][1], n3[s2][2], n_3d[s3][0], n_3d[s3][1], n_3d[s3][2]]
                    
            # 将所有负的数据和最小的一个正的数据保存到data_item中
            data_item_sub.append(data_item_positive)
            data_item.append(data_item_sub)
            processing_item +=1

            # 如果当前状态没有找到D值小于0的解，计数加一记录到csv文件并打印
            if len(data_item_sub) <= 1:
                invalid_item += 1
                writer.writerow([omega_z[s1], 
                                 n3[s2][0], n3[s2][1], n3[s2][2], 
                                 n_3d[s3][0], n_3d[s3][1], n_3d[s3][2]])
                tqdm.write(f"Invalid state found: omega_z {omega_z[s1]}, n3 {n3[s2][0]}, {n3[s2][1]}, {n3[s2][2]}, n3_d {n_3d[s3][0]}, {n_3d[s3][1]}, {n_3d[s3][2]}")
            

            # 添加进度条前缀。未找到可用控制量的状态组/已经遍历过的状态组
            state_tqdm.set_description(f"invalid/processed item {invalid_item}/{processing_item},{invalid_item/processing_item*100:.2f}%")

            processing_item_in_chunk += 1
            # 每5000个状态组合或者处理完全部数据后，就保存当前数据到pickle文件中
            if processing_item_in_chunk >= 5000 or processing_item == state_len:
                processing_item_in_chunk = 0
                n3_dataset_saver.send(data_item)
                tqdm.write(f"Saved {len(data_item)} items to file. {state_len-processing_item} items remaining!")
                if processing_item == state_len:
                    tqdm.write(f"All {len(data_item)} items saved to file.") 
                # 保存后清空当前列表以节省内存
                data_item = []
    # 全部完成后        
    n3_dataset_saver.close()
    # my_tools.save_dataset(data_item, 'found_data_set_attitude.pickle')

def find_dataset_pg():
    # 根据不同的状态pg, 寻找合适的控制量，使得最终的D<0

    # 控制量
    omega_z = np.arange(-0.2, 0.2, 0.02)
    v_zd  = np.arange(-2, 2, 0.2)
    print(f'control seed:')
    line1 = '|omega_z'.ljust(10) + '|v_zd'.ljust(10) + '|product'.ljust(10)
    print(f'{line1}')
    seed_len = len(omega_z) * len(v_zd)
    line2 = f'|{len(omega_z)}'.ljust(10) + f'|{len(v_zd)}'.ljust(10) + f'|{seed_len}'.ljust(10)
    print(f'{line2}')

    # 状态量
    p_x = np.arange(0.1, 0.4, 0.1)
    p_y = np.arange(0.1, 0.4, 0.1)
    c_z = np.arange(30, 50, 20)
    gv_z = np.arange(5, 10, 5)
    gen_n3 = interception_system.generate_n3()
    n3 = list(gen_n3)
    omega_x = np.arange(-0.2, 0.2, 0.02)
    omega_y = np.arange(-0.2, 0.2, 0.02)
    print(f'state seed:')
    line1 = '|p_x'.ljust(10) + '|p_y'.ljust(10) + '|c_z'.ljust(10)+ '|gv_z'.ljust(10)+ '|n_3'.ljust(10)+ '|omega_x'.ljust(10)+ '|omega_y'.ljust(10) + '|product'.ljust(10)
    print(f'{line1}')
    seed_len = len(p_x) * len(p_y) * len(c_z) * len(gv_z) * len(n3) * len(omega_x) * len(omega_y)
    line2 = f'|{len(p_x)}'.ljust(10) + f'|{len(p_y)}'.ljust(10)+ f'|{len(c_z)}'.ljust(10)+ f'|{len(gv_z)}'.ljust(10)+ f'|{len(n3)}'.ljust(10)+ f'|{len(omega_x)}'.ljust(10)+ f'|{len(omega_y)}'.ljust(10) + f'|{seed_len}'.ljust(10)
    print(f'{line2}')

    # 根据状态组合数量添加带进度条的循环
    state_pdt = product(range(len(p_x)), range(len(p_y)), range(len(c_z)), range(len(gv_z)), range(len(n3)),range(len(omega_x)), range(len(omega_y)))
    state_len = len(p_x) * len(p_y) * len(c_z) * len(gv_z) * len(n3) * len(omega_x) * len(omega_y)
    state_tqdm = tqdm(state_pdt, total = state_len, ncols=20)
    # 开始循环遍历
    data_item = []
    invalid_item = 0
    processing_item = 0
    for s1, s2, s3, s4, s5, s6, s7 in state_tqdm:
        data_item_sub = []
        has_positive_item = False
        for u1, u2 in product(range(len(omega_z)), range(len(v_zd))):
            D_value, D1, D2 = interception_system.cal_D_value_pg(
                p_g = np.array([p_x[s1], p_y[s2]]),
                c_z = c_z[s3],
                gv_z = gv_z[s4],
                n_3 = n3[s5],
                v_zd = v_zd[u2],
                omega = np.array([omega_x[s6], omega_y[s7], omega_z[u1]])
            )

            if D_value < 0:
                data_item_sub.append([D_value[0,0], v_zd[u2], omega_x[s6], omega_y[s7], omega_z[u1], p_x[s1], p_y[s2], c_z[s3], gv_z[s4], n3[s5]])

            if D_value >= 0:
                if not has_positive_item:
                    has_positive_item = True
                    data_item_positive = [D_value[0,0], v_zd[u2], omega_x[s6], omega_y[s7], omega_z[u1], p_x[s1], p_y[s2], c_z[s3], gv_z[s4], n3[s5]]
                else:
                    if D_value[0,0] < data_item_positive[0]:
                        data_item_positive = [D_value[0,0], v_zd[u2], omega_x[s6], omega_y[s7], omega_z[u1], p_x[s1], p_y[s2], c_z[s3], gv_z[s4], n3[s5]]
                
        # 将所有负的数据和最小的一个正的数据保存到data_item中
        data_item_sub.append(data_item_positive)
        data_item.append(data_item_sub)
        if len(data_item_sub) <= 1:
            invalid_item += 1
        processing_item +=1
        # 添加进度条前缀。未找到可用控制量的状态组/已经遍历过的状态组

        state_tqdm.set_description(f"invalid/processed item {invalid_item}/{processing_item},{invalid_item/processing_item*100:.2f}%")

    my_tools.save_dataset(data_item, 'found_data_set_pg.pickle')


COMMANDS = {
    "all": "Generate full dataset",
    "attitude": "Generate attitude control dataset only",
    "pg": "Generate pg control dataset only"
}

if __name__ == "__main__":
    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        if sys.argv[1] == 'all':
            find_dataset()
        elif sys.argv[1] == 'attitude':
            find_dataset_n3()
        elif sys.argv[1] == 'pg':
            find_dataset_pg()
        else:
            my_tools.print_help(sys.argv[0], COMMANDS)

